package com.xh.hongaiagent.agent;

import com.alibaba.cloud.ai.dashscope.api.DashScopeApi;
import com.xh.hongaiagent.agent.model.AgentStatus;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CompletableFuture;
import java.util.stream.Collectors;

@Slf4j
@Data
public abstract class BaseAgent {

      private String name;

      private String systemPrompt;

      private String nextStatePrompt;
      private AgentStatus status = AgentStatus.IDLE;

      //执行步骤控制
      private int currentStep = 0;
      //最大执行步骤
      private int maxSteps = 10;

      //llm大模型
      private ChatClient chatClient;

      //上下文记忆
      private List<Message> context = new ArrayList<>();
      public String run(String userPrompt) {
            if (status != AgentStatus.IDLE){
                  throw new RuntimeException("Cannot run agent from state "+  status);
            }
            if(StringUtils.isBlank(userPrompt)){
                  throw new RuntimeException("Cannot run agent with empty user prompt");
            }
            this.status = AgentStatus.RUNNING;
            //记录消息上下文
            context.add(new UserMessage(userPrompt));
            //开始执行步骤
            List<String> results = new ArrayList<>();
            try {
                  for(int i = 0; i < this.maxSteps && this.status != AgentStatus.FINISHED; i++){
                        int stepNum = i + 1;
                        currentStep = stepNum;
                        log.info("Executing step: {} / {}", stepNum, this.maxSteps);
                        String stepResult = step();
                        String result = "Step " + stepNum + ": " + stepResult;
                        results.add(result);
                  }
                  //检查是否超出步骤数
                  if (currentStep >= this.maxSteps) {
                        status = AgentStatus.FINISHED;
                        results.add("Terminated: Reached max step("+ this.maxSteps + ")");
                  }
                  return results.stream().collect(Collectors.joining("\n"));
            } catch (Exception e) {
                  status = AgentStatus.ERROR;
                  log.error("Executing error: ", e);
                  return "执行错误: " + e.getMessage();
            } finally {
                  clean();
            }
      }

      /**
       * 流式输出模型结果
       * @param userPrompt
       * @return
       */
      public SseEmitter runByStream(String userPrompt) {


            SseEmitter sseEmitter = new SseEmitter(300000L);

            //使用线程异步处理，避免阻塞
            CompletableFuture.runAsync(() -> {
                  try {
                        if (status != AgentStatus.IDLE){
                              sseEmitter.send("状态错误： "+ status);
                              sseEmitter.complete();
                              return;
                        }
                        if(StringUtils.isBlank(userPrompt)){
                              sseEmitter.send("输入的prompt不能为空");
                              sseEmitter.complete();
                              return;
                        }
                  } catch (Exception e) {
                        sseEmitter.completeWithError(e);
                  }
                  this.status = AgentStatus.RUNNING;
                  //记录消息上下文
                  context.add(new UserMessage(userPrompt));
                  //开始执行步骤
//                  List<String> results = new ArrayList<>();
                  try {
                        for(int i = 0; i < this.maxSteps && this.status != AgentStatus.FINISHED; i++){
                              int stepNum = i + 1;
                              currentStep = stepNum;
                              log.info("Executing step: {} / {}", stepNum, this.maxSteps);
                              String stepResult = step();
                              String result = "Step " + stepNum + ": " + stepResult +" \n";
//                              results.add(result);
                              //输入当前步骤结果到sse
                              sseEmitter.send(result);

                              // 每一步 step 执行完都要检查是否陷入循环
                              if (isStuck()) {
                                    handleStuckState();
                              }
                        }

                        //检查是否超出步骤数
                        if (currentStep >= this.maxSteps) {
                              status = AgentStatus.FINISHED;
//                              results.add("Terminated: Reached max step("+ this.maxSteps + ")");
                              sseEmitter.send("执行结束，达到最大步骤： "+ this.maxSteps);
                        }

                        sseEmitter.complete();
//                        return results.stream().collect(Collectors.joining("\n"));
                  } catch (Exception e) {
                        status = AgentStatus.ERROR;
                        log.error("Executing error: ", e);
                      try {
                          sseEmitter.send("执行错误: " + e.getMessage());
                      } catch (IOException ex) {
                          sseEmitter.completeWithError(ex);
                      }
                      sseEmitter.complete();
//                        return "执行错误: " + e.getMessage();
                  } finally {
                        clean();
                  }
            });

            //设置超时回调
            sseEmitter.onTimeout(()->{
                  this.status = AgentStatus.ERROR;
                  clean();
                  log.error(getName() + " SSE conection timeout");
            });

            //设置完成回调
            sseEmitter.onCompletion(()->{
                  if (this.status == AgentStatus.RUNNING)
                        this.status = AgentStatus.FINISHED;
                  clean();
                  log.info(getName() + " SSE conection finished");
            });

            return sseEmitter;
      }

      public abstract String step();

      protected void clean(){
            //清理资源

      }

      private int duplicateThreshold = 2;

      /**
       * 处理陷入循环的状态
       */
      protected void handleStuckState() {
            String stuckPrompt = "观察到重复响应。考虑新策略，避免重复已尝试过的无效路径。";
            this.nextStatePrompt = stuckPrompt + "\n" + (this.nextStatePrompt != null ? this.nextStatePrompt : "");
            System.out.println("Agent detected stuck state. Added prompt: " + stuckPrompt);
      }

      /**
       * 检查代理是否陷入循环
       *
       * @return 是否陷入循环
       */
      protected boolean isStuck() {
            if (context.size() < 2) {
                  return false;
            }

            Message lastMessage = context.get(context.size() - 1);
            if (lastMessage.getText() == null || lastMessage.getText().isEmpty()) {
                  return false;
            }

            // 计算重复内容出现次数
            int duplicateCount = 0;
            for (int i = context.size() - 2; i >= 0; i--) {
                  Message msg = context.get(i);
                  if (msg.getMetadata().get("role") == DashScopeApi.ChatCompletionMessage.Role.ASSISTANT &&
                          lastMessage.getText().equals(msg.getText())) {
                        duplicateCount++;
                  }
            }

            return duplicateCount >= this.duplicateThreshold;
      }

}
